'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: voceval.py
@time: 2019-08-14 16:58:18
@desc: 
'''
from jjzhk.config import ZKCFG
from jjzhk.progressbar import ProgressBar


class EvalObj:
    def __init__(self, config:ZKCFG, model):
        self.cfg = config
        self.model = model

    def calculateMAP(self, loader, output_path, detector):
        all_boxes = [[[] for _ in range(len(loader))]
                     for _ in range(len(self.cfg.BASE.CLASSINFO.keys()) + 1)]
        infos = []

        bar = ProgressBar(1, len(loader), "Detections %d")
        for i, sampler in enumerate(loader):
            bar.show(1, i)
            images, targets, info = sampler['img'], sampler['annot'], sampler['info']
            detections = self.model.get_detections(images, detector=detector)

            image_eval_boxes = self.model.get_eval_predictions(info, detections)

            for j, box in enumerate(image_eval_boxes):
                all_boxes[j][i] = box

            infos.append(info[0])
        print("calculating mAP...")
        return loader.dataset.evaluate_detections(all_boxes, output_path, infos)

